#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
from time import sleep
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper

# 配置P4表规则的通用函数，负责构建并下发表项
def configure_table_rule(switch_conn, p4_info_obj, table_name, match_fields, action_name, action_params, desc):
    try:
        # 基于P4Info构建表项
        table_entry = p4_info_obj.buildTableEntry(
            table_name=table_name,
            match_fields=match_fields,
            action_name=action_name,
            action_params=action_params
        )
        # 向交换机下发表项
        switch_conn.WriteTableEntry(table_entry)
        print(f"{desc} succeeded")
    except grpc.RpcError as e:
        print(f"{desc} failed:")
        printGrpcError(e)
        sys.exit(1)

def main():
    # 解析命令行参数，指定p4info和bmv2-json文件路径
    parser = argparse.ArgumentParser()
    parser.add_argument('--p4info', default="build/acl.p4.p4info.txt")
    parser.add_argument('--bmv2-json', default="build/acl.json")
    args = parser.parse_args()

    # 检查编译文件是否存在，若不存在提示执行make命令
    if not os.path.exists(args.p4info) or not os.path.exists(args.bmv2_json):
        print("Error: Please run 'make' first to generate compiled files")
        sys.exit(1)

    # 初始化P4InfoHelper和交换机连接
    p4_info_obj = p4runtime_lib.helper.P4InfoHelper(args.p4info)
    switch_conn = p4runtime_lib.bmv2.Bmv2SwitchConnection(
        name='s1', address='127.0.0.1:50051', device_id=0
    )

    try:
        # 获取主仲裁权并加载转发流水线配置
        switch_conn.MasterArbitrationUpdate()
        switch_conn.SetForwardingPipelineConfig(
            p4info=p4_info_obj.p4info,
            bmv2_json_file_path=args.bmv2_json
        )
        print("P4 program loaded successfully")

        # 定义主机配置列表，包含IP、MAC和端口信息
        host_configs = [
            {"ip_addr": "10.0.1.1", "mac_addr": "08:00:00:00:01:01", "port_num": 1},
            {"ip_addr": "10.0.1.2", "mac_addr": "08:00:00:00:01:02", "port_num": 2},
            {"ip_addr": "10.0.1.3", "mac_addr": "08:00:00:00:01:03", "port_num": 3},
            {"ip_addr": "10.0.1.4", "mac_addr": "08:00:00:00:01:04", "port_num": 4}
        ]
        # 为每个主机配置IPv4转发规则
        for host in host_configs:
            configure_table_rule(
                switch_conn=switch_conn,
                p4_info_obj=p4_info_obj,
                table_name="AclIngressProcessor.ipv4_forwarding_table",  # 关联P4中的IPv4转发表
                match_fields={"hdr.ipv4_hdr.dest_ip_addr": (host["ip_addr"], 32)},  # 匹配目标IP地址
                action_name="AclIngressProcessor.set_next_hop",  # 关联P4中的转发动作
                action_params={"dest_mac": host["mac_addr"], "egress_port": host["port_num"]},  # 配置目标MAC和出端口
                desc=f"Forward to host {host['ip_addr']}"
            )

        # 配置UDP 80端口过滤规则
        UDP_BLOCKED_PORT = 80
        configure_table_rule(
            switch_conn=switch_conn,
            p4_info_obj=p4_info_obj,
            table_name="AclIngressProcessor.udp_port80_filter_table",  # 关联P4中的UDP过滤表
            match_fields={"hdr.udp_hdr.udp_dest_port": UDP_BLOCKED_PORT},  # 匹配UDP目标端口80
            action_name="AclIngressProcessor.filter_udp_packet",  # 关联P4中的UDP过滤动作
            action_params={},
            desc=f"Block UDP packets to port {UDP_BLOCKED_PORT}"
        )

        # 配置MAC重写规则
        switch_mac = "08:00:00:00:01:00"
        for port in [1,2,3,4]:
            configure_table_rule(
                switch_conn=switch_conn,
                p4_info_obj=p4_info_obj,
                table_name="AclEgressProcessor.mac_rewrite_table",  # 关联P4中的MAC重写表
                match_fields={"standard_metadata.egress_port": port},  # 匹配出端口
                action_name="AclEgressProcessor.rewrite_source_mac",  # 关联P4中的MAC重写动作
                action_params={"source_mac": switch_mac},  # 配置重写后的源MAC
                desc=f"Rewrite source MAC for egress port {port}"
            )

        print("\nAll rules loaded successfully. Press Ctrl+C to exit")
        while True:
            sleep(1)
    except KeyboardInterrupt:
        print("\nExiting controller...")
    finally:
        # 关闭所有交换机连接
        ShutdownAllSwitchConnections()

if __name__ == '__main__':
    main()